from autogen_agentchat.agents import AssistantAgent
from autogen_agentchat.base import TaskResult
from autogen_agentchat.conditions import TextMentionTermination, ExternalTermination, MaxMessageTermination
from autogen_agentchat.teams import RoundRobinGroupChat
from autogen_agentchat.ui import Console
from autogen_ext.models.openai import OpenAIChatCompletionClient
model_client = OpenAIChatCompletionClient(model="qwen2.5:7b",
                                                 model_info={
                                                     "vision": False,
                                                     "function_calling": True,
                                                     "family": "Qwen3",
                                                     "structured_output": True,
                                                     "json_output": True,
                                                 },
                                                 api_key="ollama",
                                                 base_url="http://127.0.0.1:11434/v1")


# 创建主代理。
primary_agent = AssistantAgent(
    "primary",
    model_client=model_client,
    system_message="你是一个人工智能助手.",
)

# 创建评论代理。
critic_agent = AssistantAgent(
    "critic",
    model_client=model_client,
    system_message="提供建设性反馈。当你的反馈被采纳时，请回复 'APPROVE'。",
)

# 定义终止条件
text_termination = TextMentionTermination("APPROVE")
max_msg_termination = MaxMessageTermination(max_messages=10)
team = RoundRobinGroupChat([primary_agent, critic_agent], termination_condition=text_termination|max_msg_termination)

# Run the team and stream the messages to the console.
async def main() -> None:
    await Console(team.run_stream(task="今天天气如何？"))
    # Reset the team
    await team.reset()
    #上面已经触发中断，再次调用
    await Console(team.run_stream(task="写一首关于秋天的小诗."))

# NOTE: if running this inside a Python script you'll need to use asyncio.run(main()).
if __name__ == "__main__":
    import asyncio
    asyncio.run(main())